import os
from typing import TypedDict, Annotated, Any, Literal

import dotenv
import requests
import json
from langchain_community.tools import GoogleSerperRun
from langchain_community.tools.openai_dalle_image_generation import OpenAIDALLEImageGenerationTool
from langchain_community.utilities import GoogleSerperAPIWrapper
from langchain_community.utilities.dalle_image_generator import DallEAPIWrapper
from pydantic import Field, BaseModel
from langchain_core.tools import tool, Tool
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph
from langgraph.graph.message import add_messages
from langchain_core.messages import ToolMessage

dotenv.load_dotenv()


@tool
def bocha_web_search_tool(query: str, count: int = 8) -> str:
    """
    使用Bocha Web Search API进行联网搜索，返回搜索结果的字符串。

    参数:
    - query: 搜索关键词
    - count: 返回的搜索结果数量

    返回:
    - 搜索结果的字符串形式
    """
    url = 'https://api.bochaai.com/v1/web-search'
    headers = {
        'Authorization': f'Bearer {os.getenv("BOCHA_API_KEY")}',  # 请替换为你的API密钥
        'Content-Type': 'application/json'
    }
    data = {
        "query": query,
        "freshness": "noLimit",  # 搜索的时间范围，例如 "oneDay", "oneWeek", "oneMonth", "oneYear", "noLimit"
        "summary": True,  # 是否返回长文本摘要总结
        "count": count
    }

    response = requests.post(url, headers=headers, json=data)

    if response.status_code == 200:
        # 返回给大模型的格式化的搜索结果文本
        # 可以自己对博查的搜索结果进行自定义处理
        return str(response.json())
    else:
        raise Exception(f"API请求失败，状态码: {response.status_code}, 错误信息: {response.text}")


@tool
def baidu_qianfan_img_generate(query: str) -> str:
    """这是百度千帆的ai生成图片的工具。输出是图片的路径"""
    url = "https://qianfan.baidubce.com/v2/images/generations"
    payload = json.dumps({
        "prompt": query,
        "model": "irag-1.0"
    }, ensure_ascii=False)
    gaode_api_key = os.getenv("QIANFAN_API_KEY")
    headers = {
        'Content-Type': 'application/json',
        'Authorization': f'Bearer {gaode_api_key}'
    }
    response = requests.request("POST", url, headers=headers, data=payload.encode("utf-8"))
    if response.status_code == 200:
        # 返回给大模型的格式化的搜索结果文本
        # 可以自己对博查的搜索结果进行自定义处理
        return str(response.json())
    else:
        raise Exception(f"API请求失败，状态码: {response.status_code}, 错误信息: {response.text}")


class GoogleSerperArgsSchema(BaseModel):
    query: str = Field(description="执行谷歌搜索的查询语句")


class DallEArgsSchema(BaseModel):
    query: str = Field(description="输入应该是生成图像的文本提示(prompt)")


# 1. 创建博查查询工具
bocha_tool = Tool(
    name="BochaWebSearch",
    func=bocha_web_search_tool,
    description="使用Bocha Web Search API进行网络搜索",
    args_schema=GoogleSerperArgsSchema,
    api_wrapper=GoogleSerperAPIWrapper(),
)

# 2. 创建千帆文生图工具
qianfan_tool = Tool(
    name="QianfanImgGenerator",
    func=baidu_qianfan_img_generate,
    description="使用QIAN FAN irag-1.0生成图片",
    args_schema=DallEArgsSchema,
    # api_wrapper = GoogleSerperAPIWrapper(),
)

# 3.定义谷歌搜索工具
google_serper = GoogleSerperRun(
    name="google_serper",
    description=(
        "一个低成本的谷歌搜索API。"
        "当你需要回答有关时事的问题时，可以调用该工具。"
        "该工具的输入是搜索查询语句。"
    ),
    api_wrapper=GoogleSerperAPIWrapper(),
)

# 4. 定义OpenAI文生图工具
dalle = OpenAIDALLEImageGenerationTool(
    name="openai_dalle",
    api_wrapper=DallEAPIWrapper(model="dall-e-3"),
    # args_schema=DallEArgsSchema,
)


class State(TypedDict):
    """图状态的数据结构，类型为字典"""
    messages: Annotated[list, add_messages]


tools = [bocha_tool, qianfan_tool]

# 有钱请上 model="dall-e-3"
llm = ChatOpenAI(model="kimi-k2-turbo-preview")
llm_with_tools = llm.bind_tools(tools)


def chatbot(state: State) -> Any:
    """聊天机器人函数"""
    # 1. 获取状态里存储的消息列表数据并传递给LLM
    ai_message = llm_with_tools.invoke(state["messages"])
    # 2. 返回更新/生成的状态
    return {"messages": [ai_message]}


def tool_executor(state: State) -> Any:
    """工具执行节点"""
    # 1. 提取数据状态中的tool_calls
    too_calls = state["messages"][-1].tool_calls

    # 2. 根据找到的tool_calls去获取需要执行什么工具
    tools_by_name = {tool.name: tool for tool in tools}
    # 3. 执行工具得到对应的结果
    messages = []
    for tool_call in too_calls:
        tool = tools_by_name[tool_call["name"]]
        tool.invoke(tool_call["args"])
        messages.append(ToolMessage(
            tool_call_id=tool_call["id"],
            content=json.dumps(tool.invoke(tool_call["args"])),
            name=tool_call["name"],
        ))

    # 4. 将工具的执行结果作为工具消息  新到数据状态机中
    return {"messages": messages}


def route(state: State) -> Literal["tool_executor", "__end__"]:
    """通过路由检测返回的节点，返回的节点有两个，一个工具执行，一个结束节点"""
    ai_message = state["messages"][-1]
    if hasattr(ai_message, "tool_calls") and len(ai_message.tool_calls) > 0:
        return "tool_executor"
    return "__end__"


################
# 1. 创建状态图， 并使用GrathState作为状态数据
graph_builder = StateGraph(State)

# 2. 添加节点
graph_builder.add_node("llm", chatbot)
graph_builder.add_node("tool_executor", tool_executor)

# 3. 添加边
graph_builder.set_entry_point("llm")
graph_builder.add_conditional_edges("llm", route)
graph_builder.add_edge("tool_executor", "llm")

# 4. 编译图为Runnable可运行组件
graph = graph_builder.compile()

# 5. 调用图架构应用
state = graph.invoke({"messages": [("human", "2024年北京半程马拉松的前三名成绩是多少？ ")]})

for message in state["messages"]:
    print("消息类型：", message.type)
    if hasattr(message, "tool_calls") and len(message.tool_calls) > 0:
        print("工具调用参数", message.tool_calls)
    print("消息内容", message.content)
    print("==============================================")
